# -*- coding: utf-8 -*-
"""
Created on Fri Apr  8 08:56:10 2022

@author: 04566
"""
import os
import paddle
import numpy as np
from paddle.io import Dataset
from PIL import Image  # using pillow-simd for increased speed
from paddle.vision import transforms
import cv2
def read_img(fpath,width,height):
    img=cv2.imread(fpath)
    img=Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
    img=img.convert('RGB')
    '''
    with open(fpath, 'rb') as f:
        with Image.open(f) as img:
            #img=img.resize((width, height),Image.ANTIALIAS)
            return img.convert('RGB')
    '''
    return img



class dataset(Dataset):
    def __init__(self,dirpath,img_dataset_source,width,height):
        self.dirpath=dirpath
        self.img_dataset_source=img_dataset_source
        
        self.to_tensor = transforms.ToTensor()

        self.brightness = 0.2
        self.contrast = 0.2
        self.saturation = 0.2
        self.hue = 0.1
        self.resize = {}
        self.width=width
        self.height=height
        
        self.color_aug = transforms.ColorJitter()._get_param(self.brightness, self.contrast, self.saturation, self.hue)    
            
        
        self.K = np.array([[0.58, 0, 0.5, 0],
                           [0, 1.92, 0.5, 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]]).astype('float32')
        self.K[0,:]*=self.width
        self.K[1,:]*=self.height
        self.T = np.array([[1,0,0,-0.1],
                           [0,1,0,0],
                           [0,0,1,0],
                           [0,0,0,1]]).astype('float32')
        self.temp_datas=[]
        with open(os.path.join(self.dirpath,'{}_files.txt'.format(self.img_dataset_source)),'r',encoding='utf-8') as filenames:
            self.temp_datas=filenames.readlines()
    
    def __getitem__(self, idx):
        temp=self.temp_datas[idx].split()[0]
        temp=os.path.join('/home/aistudio/data/kitti',temp)           
        img_name=self.temp_datas[idx].split()[1]
        place=self.temp_datas[idx].split()[2]
        inputs={}
        if place=='l':
            img_path=os.path.join(temp,'image_02','data',str(img_name).zfill(10)+'.jpg')
            l=read_img(img_path,self.width,self.height)
            img_path=os.path.join(temp,'image_03','data',str(img_name).zfill(10)+'.jpg')
            r=read_img(img_path,self.width,self.height)
        else:
            img_path=os.path.join(temp,'image_03','data',str(img_name).zfill(10)+'.jpg')
            l=read_img(img_path,self.width,self.height)
            img_path=os.path.join(temp,'image_02','data',str(img_name).zfill(10)+'.jpg')
            r=read_img(img_path,self.width,self.height)
        for i in range(4):
            s = 2 ** i
            l=l.resize((self.width // s,self.height // s),Image.ANTIALIAS)
            #l=cv2.resize(l,(self.width // s,self.height // s),interpolation=cv2.INTER_AREA)
            inputs[('img_raw_l',i)]=self.to_tensor(l)
            r=r.resize((self.width // s,self.height // s),Image.ANTIALIAS)
            #r=cv2.resize(r,(self.width // s,self.height // s),interpolation=cv2.INTER_AREA)
            inputs[('img_raw_r',i)]=self.to_tensor(r)
        
        if place=='r':
            self.T[0,3]*=-1
        inputs['T']=paddle.to_tensor(self.T)
        inputs['K']=paddle.to_tensor(self.K)
        inputs['inv_K']=paddle.to_tensor(np.linalg.pinv(self.K))
        return inputs

    def __len__(self):
        return len(self.temp_datas)
